package com.smh;

import org.junit.jupiter.api.Test;

import java.util.Arrays;
import java.util.TreeMap;

/**
 * @author shiminghui
 * @date 2025/3/4 19:39
 * @description: TODO
 */
public class _036_B树 {


    @Test
    public void test1() {
        BTree<Integer> bTree = new BTree<>();
        bTree.put(1, 1);
        bTree.put(2, 2);
        bTree.put(3, 3);
        bTree.put(4, 4);
        bTree.put(5, 5);
        bTree.put(6, 6);
        bTree.put(7, 7);
        bTree.put(8, 8);
        bTree.put(9, 9);
    }


    @Test
    public void test2() {
        BTree<Integer> bTree = new BTree<>(3);
        bTree.put(1, 1);
        bTree.put(2, 2);
        bTree.put(3, 3);
        bTree.put(4, 4);
        bTree.put(5, 5);
        bTree.put(6, 6);
        bTree.put(7, 7);
        bTree.put(8, 8);
        bTree.put(9, 9);
        bTree.put(10, 10);
        bTree.put(11, 11);
        bTree.put(12, 12);
        bTree.put(13, 13);
        bTree.put(14, 14);
        bTree.put(15, 15);

        System.out.println(bTree.remove(2));
        // 删除所有的值
//        for (int i = 0; i < 15; i++) {
//            bTree.remove(i + 1);
//        }
    }


    private static class BTree<T> {

        private final int M; // 节点的最小度数
        private final int MIN_KEY_NUM;
        private final int MAX_KEY_NUM;

        private BTreeNode<T> root = null;

        public BTree() {
            this(2);
        }

        public BTree(int M) {
            this.M = M;
            root = new BTreeNode<>(M);
            MIN_KEY_NUM = M - 1;
            MAX_KEY_NUM = M * 2 - 1;
        }

        /**
         * 根据key删除节点
         *
         * @param key
         * @return
         */
        public T remove(int key) {
            return doRemove(root, key, null, -1);
        }

        /**
         * @param node   当前节点
         * @param key    删除的key值
         * @param parent 当前节点的父节点
         * @param index  当前节点在父节点孩子中的下标
         * @return
         */
        private T doRemove(BTreeNode<T> node, int key, BTreeNode<T> parent, int index) {
            if (node == null) { // 节点为空，直接返回
                return null;
            }
            T result = null;
            int p = 0;
            for (int i = 0; i < node.keyNum; i++) {
                if (node.keys[i] >= key)
                    break;
                p++;
            }
            // 找到了key
            if (findNode(node, key, p)) {
                // 是叶子节点
                if (node.isLeaf) {
                    result = node.removeKeyAndValue(p);
                }
                // 不是叶子节点
                else {
                    //todo
                    // 寻找其后继节点，然后将后继节点的值替换当前节点的值，然后删除后继节点
                    BTreeNode<T> child = node.children[p + 1];
                    while (!child.isLeaf) {
                        child = child.children[0];
                    }
                    // 此时的child就是其后继节点
                    // 交换key和value值
                    node.keys[p] = child.keys[0];
                    node.values[p] = child.values[0];
                    // 删除后继节点
                    result = doRemove(node.children[p + 1], child.keys[0], node, p + 1);
                }
            }
            // 当前节点没找到key
            else {
                // 是叶子节点
                if (node.isLeaf) {
                    return null;
                }
                // 不是叶子节点
                else {
                    result = doRemove(node.children[p], key, node, p);
                }
            }
            // 调整平衡
            if (node.keyNum < MIN_KEY_NUM) {
                // todo
                balance(node, parent, index);
            }
            return result;
        }

        /**
         * 平衡节点
         *
         * @param node   平衡节点
         * @param parent 平衡节点的父节点
         * @param index  平衡节点在父节点孩子中的下标
         */
        private void balance(BTreeNode<T> node, BTreeNode<T> parent, int index) {
            if (node == null) {
                return;
            }
            if (parent == null) { // 根节点,不用调整
                if (node.keyNum == 0 && node.children[0] != null)
                    root = node.children[0];
                return;
            }
            // 平衡节点的左右兄弟节点
            BTreeNode<T> leftChild = parent.childLeftChild(index);
            BTreeNode<T> rightChild = parent.childRightChild(index);

            if (leftChild != null && leftChild.keyNum > MIN_KEY_NUM) { // 左兄弟节点不为空，并且大于最小度数，则从左兄弟节点借
                // 先从父亲节点借一个key和value
                node.insertKeyAndValue(parent.keys[index - 1], parent.values[index - 1], 0);
                // 从左兄弟节点删除一个key和value，放到父亲节点
                parent.keys[index - 1] = leftChild.keys[leftChild.keyNum - 1];
                parent.values[index - 1] = leftChild.values[leftChild.keyNum - 1];
                leftChild.removeKeyAndValue(leftChild.keyNum - 1);
                BTreeNode<T> removeMostRightChild = leftChild.removeMostRightChild();
                // 当左兄弟节点不是叶子节点，则需将其最右侧的孩子放到平衡节点的最左侧孩子处
                if (!leftChild.isLeaf) {
                    node.insertChild(removeMostRightChild, 0);
                }
            } else if (rightChild != null && rightChild.keyNum > MIN_KEY_NUM) { // 右兄弟节点不为空，并且大于最小度数，则从右兄弟节点借
                node.insertKeyAndValue(parent.keys[index], parent.values[index], node.keyNum);
                parent.keys[index] = rightChild.keys[0];
                parent.values[index] = rightChild.values[0];
                rightChild.removeKeyAndValue(0);
                BTreeNode<T> removeMostLeftChild = rightChild.removeMostLeftChild();
                if (!rightChild.isLeaf) {
                    node.insertChild(removeMostLeftChild, node.keyNum); // 这里keyNum 不用加1，因为之前insertKeyAndValue方法已经加1了，这里的keyNum就是孩子的数量。
                }
            } else if (leftChild != null) { // 左右兄弟为空或者都小于最小度数
                // 这时候要进行合并了
                leftChild.insertKeyAndValue(parent.keys[index - 1], parent.values[index - 1], leftChild.keyNum);
                parent.removeKeyAndValue(index - 1);
                parent.removeChild(index);
                node.moveToTarget(leftChild);
            } else if (rightChild != null) { // 右兄弟不为空，并且key不大于最小度数
                node.insertKeyAndValue(parent.keys[index], parent.values[index], node.keyNum);
                parent.removeKeyAndValue(index);
                parent.removeChild(index + 1);
                rightChild.moveToTarget(node);
            }
        }

        /**
         * 查询key是否存在
         *
         * @param key
         * @return
         */
        public boolean contains(int key) {
            BTreeNode<T> node = root.get(key);
            return node != null;
        }

        /**
         * 插入key和value
         *
         * @param key
         * @param value
         */
        public void put(int key, T value) {
            doPut(root, key, value, null, -1);
        }

        /**
         * 递归查找，并插入key和value
         *
         * @param node
         * @param key
         * @param value
         * @param parent     父节点
         * @param childIndex 在父节点孩子中的下标
         */
        private void doPut(BTreeNode<T> node, int key, T value, BTreeNode<T> parent, int childIndex) {
            int p = 0;
            for (int i = 0; i < node.keyNum; i++) {
                if (node.keys[i] >= key)
                    break;
                p++;
            }
            if (findNode(node, key, p)) {
                // 存在key，更新value
                node.values[p] = value;
            } else {
                // 是叶子节点
                if (node.isLeaf) {
                    // 当前节点的key数量小于最大度数，直接插入
                    node.insertKeyAndValue(key, value, p);
                }
                // 不是叶子节点
                else {
                    // 继续递归
                    doPut(node.children[p], key, value, node, p);
                }
                if (node.keyNum == MAX_KEY_NUM) {
                    splitNode(node, parent, childIndex);
                }
            }
        }

        /**
         * 分裂节点
         *
         * @param node   分裂节点
         * @param parent 父节点
         * @param index  是父节点孩子的下标
         */
        private void splitNode(BTreeNode<T> node, BTreeNode<T> parent, int index) {
            // 父节点为空，创建一个根节点
            if (parent == null) {
                BTreeNode<T> newRoot = new BTreeNode<>(M);
                newRoot.isLeaf = false;
                newRoot.insertChild(node, 0);
                root = newRoot;
                parent = root;
                index = 0;
            }

            BTreeNode<T> newNode = new BTreeNode<>(M);
            newNode.isLeaf = node.isLeaf; // 新节点是否为叶子节点跟旧节点的子节点是否是叶子节点相同
            // key和value拷贝
            System.arraycopy(node.keys, M, newNode.keys, 0, M - 1);
            System.arraycopy(node.values, M, newNode.values, 0, M - 1);
            if (!node.isLeaf) {
                // 旧节点不是叶子节点，需要拷贝旧节点的孩子节点
                System.arraycopy(node.children, M, newNode.children, 0, M);
            }

            // 向父节点插入key和value，以及新孩子
            parent.insertKeyAndValue(node.keys[M - 1], node.values[M - 1], index); //将中间的key插入到父亲节点中
            parent.insertChild(newNode, index + 1);//父亲节点多一个孩子，新孩子的索引为index+1；旧节点的index不变
            // 旧节点的key和value更新
            node.keyNum = M - 1; // 旧节点的key变为M-1，即最小的度数-1
            newNode.keyNum = M - 1; // 新节点的key变为M-1；


        }

        /**
         * 查找key是否在当前节点中
         *
         * @param node
         * @param key
         * @param index
         * @return
         */
        private boolean findNode(BTreeNode<T> node, int key, int index) {
            return node.keys[index] == key;
        }


        private static class BTreeNode<T> {
            private BTreeNode<T>[] children; // 当前节点的子节点
            private T[] values; // 当前节点的值
            private int[] keys; // 当前节点的key
            private int keyNum; // 当前节点的key数量
            private boolean isLeaf = true; // 是否是叶子节点
            private int size; // 当前节点的size，也是最小度数

            BTreeNode(int size) {
                this.size = size;
                keys = new int[size * 2 - 1];
                values = (T[]) new Object[size * 2 - 1];
                children = (BTreeNode<T>[]) java.lang.reflect.Array.newInstance(BTreeNode.class, size * 2);
                // 将keys的每个位置设置为Integer.MAX_VALUE，用于二分查找
//                Arrays.fill(keys, Integer.MAX_VALUE);
            }

            BTreeNode() {
                this(2); // 默认最小度数为2
            }


            @Override
            public String toString() {
                return Arrays.toString(Arrays.copyOfRange(keys, 0, keyNum));
            }

            /**
             * 根据key查询key所在节点
             *
             * @param key
             * @return
             */
            private BTreeNode<T> get(int key) {
                int p = 0;
                for (int i = 0; i < keyNum; i++) {
                    if (keys[i] >= key)
                        break;
                    p++;
                }
                if (keys[p] == key) {
                    return this;
                } else {
                    if (isLeaf) {
                        return null;
                    }
                    return children[p].get(key);
                }
            }

            private void insertKeyAndValue(int key, T value, int index) {
                // 插入key
                System.arraycopy(keys, index, keys, index + 1, keyNum - index);
                keys[index] = key;
                // 插入value
                System.arraycopy(values, index, values, index + 1, keyNum - index);
                values[index] = value;
                keyNum++;
            }

            private void insertChild(BTreeNode<T> child, int index) {
                System.arraycopy(children, index, children, index + 1, keyNum - index);
                children[index] = child;
            }

            /**
             * 删除指定index处的key和value
             *
             * @param index
             */
            private T removeKeyAndValue(int index) {
                T temp = values[index];
                System.arraycopy(keys, index + 1, keys, index, keyNum - index - 1);
                System.arraycopy(values, index + 1, values, index, keyNum - index - 1);
                keyNum--;
                return temp;
            }

            /**
             * 移除指定索引的孩子节点
             *
             * @param index
             * @return
             */
            private BTreeNode<T> removeChild(int index) {
                System.arraycopy(children, index + 1, children, index, keyNum - index + 1);
                return children[index];
            }

            private T removeMostLeftKeyAndValue() {
                return removeKeyAndValue(0);
            }

            private T removeMostRightKeyAndValue() {
                return removeKeyAndValue(keyNum - 1);
            }

            private BTreeNode<T> removeMostLeftChild() {
                return removeChild(0);
            }

            private BTreeNode<T> removeMostRightChild() {
                return removeChild(keyNum);
            }

            private BTreeNode<T> childLeftChild(int index) {
                return index <= 0 ? null : children[index - 1];
            }

            private BTreeNode<T> childRightChild(int index) {
                return index >= keyNum ? null : children[index + 1];
            }

            private void moveToTarget(BTreeNode<T> target) {
                int index = target.keyNum;
                // 拷贝key和值
                System.arraycopy(keys, 0, target.keys, index, keyNum);
                System.arraycopy(values, 0, target.values, index, keyNum);
                // 拷贝孩子节点
                if (!isLeaf) {
                    System.arraycopy(children, 0, target.children, index + 1, keyNum + 1);
                }
            }
        }
    }

}
